'''
@Descripttion: this method usually consumes 95.34 s(slow)
@Author: Chen Chenxi
@Date: 2019-10-25 14:38:11
@LastEditTime: 2019-11-07 11:35:45
'''
import os
import time
import numpy as np
import cv2
from glob import glob


def cal_for_frames(video_path):
    frames = glob(os.path.join(video_path, '*.jpg'))
    frames.sort()

    frames = frames[:10]
    flow = []
    prev = cv2.imread(frames[0])
    prev = cv2.cvtColor(prev, cv2.COLOR_BGR2GRAY)
    for i, frame_curr in enumerate(frames):
        curr = cv2.imread(frame_curr)
        curr = cv2.cvtColor(curr, cv2.COLOR_BGR2GRAY)
        tmp_flow = compute_TVL1(prev, curr)
        # print(tmp_flow)
        flow.append(tmp_flow)
        prev = curr
    
    e = time.time()
 
    return flow
 
 
def compute_TVL1(prev, curr, bound=15):
    TVL1=cv2.optflow.DualTVL1OpticalFlow_create()
    flow = TVL1.calc(prev, curr, None)
    assert flow.dtype == np.float32
 
    flow = (flow + bound) * (256.0 / (2 * bound))
    flow = np.round(flow).astype(int)
    flow[flow >= 256] = 256
    flow[flow <= 0] = 0
 
    return flow
 
 
def save_flow(video_flows, flow_path):
    for i, flow in enumerate(video_flows):
        if not os.path.exists(os.path.join(flow_path,'u')):
            os.mkdir(os.path.join(flow_path,'u'))
        cv2.imwrite(os.path.join(flow_path,'u', "frame{:06d}.jpg".format((i+1))),flow[:, :, 0])
        if not os.path.exists(os.path.join(flow_path, 'v')):
            os.mkdir(os.path.join(flow_path, 'v'))
        cv2.imwrite(os.path.join(flow_path,'v', "frame{:06d}.jpg".format((i+1))),flow[:, :, 1])
 
 
def extract_flow(video_path, flow_path):
    flow = cal_for_frames(video_path)
    save_flow(flow, flow_path)

    return
 
 
if __name__ == '__main__':
    # #提取好的帧的路径
    # video_path='C:/Users/65152/Desktop/rec_test/sample_img'
    # #存储光流的路径
    # save_path='C:/Users/65152/Desktop/rec_test/flow'

    #提取好的帧的路径
    video_path='C:/Users/65152/Desktop/output'
    #存储光流的路径
    save_path='C:/Users/65152/Desktop/UCF101-opticalflow/ApplyEyeMakeup'
    extract_flow(video_path,save_path)
    e = time.time()
